package com.dtease.agent.sse.server.service;

import com.dtease.agent.sse.server.model.EventMessage;
import com.dtease.agent.sse.server.util.Chunker;
import lombok.Getter;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.core.publisher.Sinks;
import reactor.core.scheduler.Scheduler;
import reactor.core.scheduler.Schedulers;

import javax.annotation.PreDestroy;
import java.time.Duration;
import java.util.Objects;
import java.util.Queue;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * @description SSE服务类
 * @author chentl
 * @class SseService
 * @package com.dtease.agent.sse.server.service
 * @since 2025/5/4 08:36
 */
@Service
public class SseService {

    private final ConcurrentHashMap<String, Sinks.Many<EventMessage>> sinksMap = new ConcurrentHashMap<>();
    @Getter
    private final Scheduler sharedScheduler = Schedulers.newBoundedElastic(
            10,
            1000,
            "shared-processor"
    );

    // 存储每个客户端的消息队列
    private final ConcurrentHashMap<String, Queue<EventMessage>> messageQueues = new ConcurrentHashMap<>();
    // 存储每个客户端的处理标志，用于控制消息处理流程
    private final ConcurrentHashMap<String, AtomicBoolean> processingFlags = new ConcurrentHashMap<>();
    // 消息队列的最大容量
    private static final int MAX_QUEUE_SIZE = 1000;

    /**
     * 获取指定客户端的事件流
     * @param sessionId 客户端ID
     * @return 事件消息的Flux流
     */
    public Flux<EventMessage> getEventStream(String sessionId) {
        return sinksMap.computeIfAbsent(sessionId, id -> {
            // 若客户端对应的消息发送器不存在，则创建一个
            Sinks.Many<EventMessage> sink = Sinks.many().multicast().onBackpressureBuffer();
            // 初始化客户端相关资源
            initClientResources(id);
            return sink;
        }).asFlux().doOnCancel(() -> cleanupClientResources(sessionId));
    }

    /**
     * 将消息加入指定客户端的消息队列
     * @param sessionId 客户端ID
     * @param message 事件消息
     */
    public void enqueueMessage(String sessionId, EventMessage message) {
        // 获取客户端的消息队列，若不存在则创建
        Queue<EventMessage> queue = messageQueues.computeIfAbsent(sessionId,
                k -> new ConcurrentLinkedQueue<>());

        // 若队列未满，则加入消息并触发处理
        if (queue.size() < MAX_QUEUE_SIZE) {
            queue.offer(message);
            triggerProcessing(sessionId);
        } else {
            // 若队列已满，发送溢出错误消息
            sendOverflowError(sessionId);
        }
    }

    /**
     * 触发指定客户端的消息处理
     * @param sessionId 客户端ID
     */
    private void triggerProcessing(String sessionId) {
        if(Objects.isNull(processingFlags.get(sessionId))){
            // 如果处理标志不存在，则初始化
            processingFlags.put(sessionId, new AtomicBoolean(false));
        }

        if (!processingFlags.get(sessionId).compareAndSet(false, true)) {
            return;
        }

        // 在共享调度器上异步执行消息处理
        sharedScheduler.schedule(() -> {
            processQueue(sessionId);
        });
    }

    /**
     * 处理指定客户端的消息队列
     * @param sessionId 客户端ID
     */
    private void processQueue(String sessionId) {
        // 获取客户端的消息队列
        Queue<EventMessage> queue = messageQueues.get(sessionId);
        if (queue == null) {
            return;
        }

        // 当队列不为空且处理标志为true时，持续处理消息
        AtomicBoolean isProcessing = processingFlags.get(sessionId);
        while (!queue.isEmpty() && Objects.nonNull(isProcessing) && isProcessing.get()) {
            EventMessage message = queue.poll();
            processMessage(sessionId, message);
        }
        // 处理完成，将处理标志设为false
        isProcessing.set(false);
    }

    /**
     * 处理单条消息，依次发送开始事件、消息块、完成事件
     * @param sessionId 客户端ID
     * @param message 事件消息
     */
    private void processMessage(String sessionId, EventMessage message) {
        Flux.concat(
                        // 发送开始事件
                        sendStartEvent(sessionId, message).flux(),
                        // 发送消息块
                        sendChunks(sessionId, message),
                        // 发送完成事件
                        sendCompleteEvent(sessionId, message).flux()
                ).subscribeOn(sharedScheduler)
                .blockLast();
    }

    /**
     * 发送错误消息给指定客户端
     * @param sessionId 客户端ID
     * @param ex 异常信息
     */
    private void sendError(String sessionId, Throwable ex) {
        sendDirect(sessionId, EventMessage.builder()
                .type("error")
                .content("处理消息时发生错误: " + ex.getMessage())
                .title("系统错误")
                .timestamp(System.currentTimeMillis())
                .build());
    }

    /**
     * 发送消息开始事件
     * @param sessionId 客户端ID
     * @param message 事件消息
     * @return Mono<Void>
     */
    private Mono<Void> sendStartEvent(String sessionId, EventMessage message) {
        return Mono.fromRunnable(() ->
                sendDirect(sessionId, EventMessage.builder()
                        .type(message.getType() + "_start")
                        .content("")
                        .title(message.getTitle())
                        .timestamp(System.currentTimeMillis())
                        .build())
        ).then();
    }

    /**
     * 分块发送消息内容
     * @param sessionId 客户端ID
     * @param message 事件消息
     * @return Flux<Void>
     */
    private Flux<Void> sendChunks(String sessionId, EventMessage message) {
        // 发送分块消息
        Flux<Void> chunkedMessages = Flux.fromIterable(Chunker.chunkByLength(message.getContent(), 12))
                .delayElements(Duration.ofMillis(100))
                .concatMap(chunk ->
                        Mono.fromRunnable(() ->
                                sendDirect(sessionId, EventMessage.builder()
                                        .type(message.getType() + "_chunk")
                                        .content(chunk)
                                        .title(message.getTitle())
                                        .timestamp(System.currentTimeMillis())
                                        .build())
                        ).then()
                );

        return chunkedMessages;
    }

    /**
     * 发送消息完成事件
     * @param sessionId 客户端ID
     * @param message 事件消息
     * @return Mono<Void>
     */
    private Mono<Void> sendCompleteEvent(String sessionId, EventMessage message) {
        return Mono.fromRunnable(() ->
                sendDirect(sessionId, EventMessage.builder()
                        .type(message.getType() + "_complete")
                        .content("")
                        .title(message.getTitle())
                        .timestamp(System.currentTimeMillis())
                        .build())
        ).then();
    }

    /**
     * 直接向指定客户端发送消息
     * @param sessionId 客户端ID
     * @param message 事件消息
     */
    private void sendDirect(String sessionId, EventMessage message) {
        Sinks.Many<EventMessage> sink = sinksMap.get(sessionId);
        if (sink != null) {
            sink.tryEmitNext(message);
        }
    }

    /**
     * 初始化指定客户端的资源
     * @param sessionId 客户端ID
     */
    private void initClientResources(String sessionId) {
        processingFlags.putIfAbsent(sessionId, new AtomicBoolean(false));
        messageQueues.putIfAbsent(sessionId, new ConcurrentLinkedQueue<>());
    }

    /**
     * 清理指定客户端的资源
     * @param sessionId 客户端ID
     */
    private void cleanupClientResources(String sessionId) {
        sinksMap.remove(sessionId);
        messageQueues.remove(sessionId);
        processingFlags.remove(sessionId);
    }

    /**
     * 发送队列溢出错误消息给指定客户端
     * @param sessionId 客户端ID
     */
    private void sendOverflowError(String sessionId) {
        sendDirect(sessionId, EventMessage.builder()
                .type("error")
                .content("队列已满，最多允许" + MAX_QUEUE_SIZE + "条消息排队")
                .title("系统错误")
                .timestamp(System.currentTimeMillis())
                .build());
    }

    /**
     * 服务销毁时关闭共享调度器
     */
    @PreDestroy
    public void shutdown() {
        sharedScheduler.dispose();
    }
}